package net.cyue.ort;

import ai.onnxruntime.OnnxJavaType;

public enum DataType {
    INT8,
    INT16,
    INT32,
    INT64,
    FLOAT16,
    BFLOAT16,
    FLOAT32,
    FLOAT64,
    STRING,
    UNKNOWN;

    public static DataType fromString(String type) {
        switch (type) {
            case "int8": return INT8;
            case "int16": return INT16;
            case "int32": return INT32;
            case "int64": return INT64;
            case "float16": return FLOAT16;
            case "bfloat16": return BFLOAT16;
            case "float": return FLOAT32;
            case "float32": return FLOAT32;
            case "float64": return FLOAT64;
            case "string": return STRING;
            default: return UNKNOWN;
        }
    }

    public static OnnxJavaType toOnnxJavaType(DataType dataType) {
        switch (dataType) {
            case INT8: return OnnxJavaType.INT8;
            case INT16: return OnnxJavaType.INT16;
            case INT32: return OnnxJavaType.INT32;
            case INT64: return OnnxJavaType.INT64;
            case FLOAT16: return OnnxJavaType.FLOAT16;
            case BFLOAT16: return OnnxJavaType.BFLOAT16;
            case FLOAT32: return OnnxJavaType.FLOAT;
            case FLOAT64: return OnnxJavaType.DOUBLE;
            case STRING: return OnnxJavaType.STRING;
            default: return OnnxJavaType.FLOAT;
        }
    }
    public static int getElementSize(DataType dataType) {
        switch (dataType) {
            case INT8: return 1;
            case INT16: return 2;
            case INT32: return 4;
            case INT64: return 8;
            case FLOAT16: return 2;
            case BFLOAT16: return 2;
            case FLOAT32: return 4;
            case FLOAT64: return 8;
            default: return 4;
        }
    }

    public OnnxJavaType toOnnxJavaType() {
        return toOnnxJavaType(this);
    }
    public int getElementSize() {
        return getElementSize(this);
    }
}
